﻿#include "CFMP4Appender.h"
#include <string>
#include <fstream>

static unsigned int read32(const char* data) {
	return ((static_cast<unsigned int>(data[0]) << 24 & 0xff000000)
		| (static_cast<unsigned int>(data[1]) << 16 & 0xff0000)
		| (static_cast<unsigned int>(data[2]) << 8 & 0xff00)
		| (static_cast<unsigned int>(data[3]) & 0xff));
}

static unsigned long long read64(const char* data) {
	unsigned int high, low;
	high = read32(data);
	low = read32(data + sizeof(unsigned int));
	return ((static_cast<unsigned long long>(high) << 32 & 0xffffffff00000000) | low);
}

static bool parseBox(char* data, Box* box)
{
	char fourBytes[4] = { 0 };
	memcpy(fourBytes, data, sizeof(fourBytes));
	memcpy(&box->size, fourBytes, sizeof(int));
	std::reverse((char*)&box->size, (char*)&box->size + sizeof(int));
	memcpy(box->type, data + sizeof(fourBytes), sizeof(int));
	box->type[4] = '\0';
	box->buffer = data;
	return true;
}

static bool parseBox(std::fstream* file, Box* box, bool withData) {
	char fourBytes[4] = { 0 };
	if (!file->read(fourBytes, sizeof(fourBytes))) {
		return false;
	}
	//memcpy(&box->size, fourBytes, sizeof(int));
	//std::reverse((char*)&box->size, (char*)&box->size +sizeof(int));
	box->size = read32(fourBytes);
	if (!file->read(box->type, sizeof(int))) {
		return false;
	}
	box->type[4] = '\0';
	if (withData) {
		box->buffer = new char[box->size];
		memcpy(box->buffer, fourBytes, sizeof(fourBytes));
		memcpy(box->buffer + sizeof(fourBytes), box->type, sizeof(int));
		if (!file->read(box->buffer + 8, box->size - 8)) {
			return false;
		}
	}
	else {
		file->seekg(box->size - 8, std::ios::cur);
	}
	return true;
}

CFMP4Appender::CFMP4Appender()
{
}

CFMP4Appender::~CFMP4Appender()
{
}

int CFMP4Appender::Append(const char* fmp4AnteriorSegmentFile, const char* fmp4PosteriorSegmentFile)
{
	int ret = -1;
	std::fstream* fpAnterior = nullptr;
	std::fstream* fpPosterior = nullptr;
	MoofBox* moofBoxLast = nullptr;
	MfraBox* mfraBoxLast = nullptr;
	Box* box = nullptr;
	unsigned long long track1BaseMediaDecodeTime = 0;
	unsigned long long track2BaseMediaDecodeTime = 0;
	unsigned long long track1BaseMediaDecodeTimeOfAnteriorFile = 0;
	unsigned long long track2BaseMediaDecodeTimeOfAnteriorFile = 0;
	do 
	{
		int sequenceNumber = 0;
		moofBoxLast = findLastMoofBox(fmp4AnteriorSegmentFile);
		if (nullptr == moofBoxLast || 0 == moofBoxLast->moofOffset) {
			ret = -2;
			printf("invalid moof box\n");
			break;
		}
		track1BaseMediaDecodeTime = moofBoxLast->trafBox1.trunBox.eachSampleDuration
			* moofBoxLast->trafBox1.trunBox.sampleCount
			+ moofBoxLast->trafBox1.tfdtBox.baseMediaDecodeTime;
		track2BaseMediaDecodeTime = moofBoxLast->trafBox2.trunBox.eachSampleDuration
			* moofBoxLast->trafBox2.trunBox.sampleCount
			+ moofBoxLast->trafBox2.tfdtBox.baseMediaDecodeTime;
		track1BaseMediaDecodeTimeOfAnteriorFile = track1BaseMediaDecodeTime;
		track2BaseMediaDecodeTimeOfAnteriorFile = track2BaseMediaDecodeTime;
		//mfraBoxLast = getMfraBox(fmp4AnteriorSegmentFile);
		//if (nullptr == mfraBoxLast) {
		//	break;
		//}
		//printf("Anterior file valid length:%lld\n", moofBoxLast->moofOffset);
		//int result = ftruncateFile(fmp4AnteriorSegmentFile, invalidLength);
		//if (result != invalidLength)
		//{
		//	break;
		//}
		fpAnterior = new std::fstream(fmp4AnteriorSegmentFile, std::ios::binary | std::ios::in | std::ios::out);
		if (!fpAnterior->is_open()) {
			ret = -2;
			break;
		}
		fpAnterior->seekp(moofBoxLast->moofOffset);
		fpPosterior = new std::fstream(fmp4PosteriorSegmentFile, std::ios::binary | std::ios::in);
		if (!fpPosterior->is_open()) {
			ret = -3;
			break;
		}
		fpPosterior->seekp(0, std::ios::end);
		unsigned long long fileSize = fpPosterior->tellp();
		fpPosterior->seekp(0, std::ios::beg);
		unsigned long long offset = 0;
		unsigned long long lastMoofOffset = 0;
		unsigned long long readSize = 0;
		do 
		{
			box = new Box;
			if (!parseBox(fpPosterior, box, true)) {
				ret = -4;
				break;
			}
			readSize += box->size;
			if (!(strcmp(box->type, "moof") == 0
				|| strcmp(box->type, "mdat") == 0
				|| strcmp(box->type, "mfra") == 0)) {
				delete box;
				box = new Box;
				continue;
			}
			offset += box->size;
			if (strcmp(box->type, "moof") == 0) {
				MoofBox moofBox(box);
				if (!moofBox.parse()) {
					break;
				}
				moofBox.mfhdBox.modifySequenceNumber(moofBoxLast->mfhdBox.sequenceNumber);
				lastMoofOffset = moofBoxLast->moofOffset + offset - box->size;
				moofBox.trafBox1.tfhdBox.modifyBaseDataOffset(lastMoofOffset);
				moofBox.trafBox2.tfhdBox.modifyBaseDataOffset(lastMoofOffset);
				track1BaseMediaDecodeTime = track1BaseMediaDecodeTimeOfAnteriorFile
					+ moofBox.trafBox1.tfdtBox.baseMediaDecodeTime;
				moofBox.trafBox1.tfdtBox.modifyBaseMediaDecodeTime(track1BaseMediaDecodeTimeOfAnteriorFile);
				track2BaseMediaDecodeTime = track2BaseMediaDecodeTimeOfAnteriorFile
					+ moofBox.trafBox2.tfdtBox.baseMediaDecodeTime;
				moofBox.trafBox2.tfdtBox.modifyBaseMediaDecodeTime(track2BaseMediaDecodeTimeOfAnteriorFile);
				fpAnterior->write(box->buffer, box->size);
			}
			else if (strcmp(box->type, "mfra") == 0) {
				MfraBox mfraBox(box);
				if (!mfraBox.parse()) {
					break;
				}
				mfraBox.tfraBox1.modifyTime(track1BaseMediaDecodeTime);
				mfraBox.tfraBox2.modifyTime(track2BaseMediaDecodeTime);
				mfraBox.tfraBox1.modifyMoofOffset(lastMoofOffset);
				mfraBox.tfraBox2.modifyMoofOffset(lastMoofOffset);
				fpAnterior->write(box->buffer, box->size);
			}
			else
			{
				if (!fpAnterior->write(box->buffer, box->size)) {
					ret = -5;
					break;
				}
				if (readSize == fileSize) {
					break;
				}
			}
		} while (true);
		ret = 0;
	} while (false);
	if (nullptr != fpAnterior) {
		fpAnterior->close();
		delete fpAnterior;
		fpAnterior = nullptr;
	}
	if (nullptr != fpPosterior) {
		fpPosterior->close();
		delete fpPosterior;
		fpPosterior = nullptr;
	}
	if (nullptr != moofBoxLast) {
		delete moofBoxLast;
		moofBoxLast = nullptr;
	}
	if (nullptr != mfraBoxLast) {
		delete mfraBoxLast;
		mfraBoxLast = nullptr;
	}
	if (nullptr != box) {
		delete box;
	}
	return ret;
}

MoofBox* CFMP4Appender::findLastMoofBox(const char* filePath) {
	unsigned long long offset = 0;
	std::fstream* fp = nullptr;
	MoofBox* moofBox = nullptr;
	TrafBox lastTrafBox1;
	TrafBox lastTrafBox2;
	do
	{
		fp = new std::fstream(filePath, std::ios::binary | std::ios::in);
		if (!fp->is_open()) {
			printf("open file fail: %s\n", filePath);
			break;
		}

		Box* box = nullptr;
		do 
		{
			if (nullptr != box) {
				delete box;
				box = nullptr;
			}
			box = new Box;
			if (!parseBox(fp, box, true)) {
				break;
			}
			offset += box->size;
			if (strcmp(box->type, "moof") == 0) {
				if (nullptr != moofBox) {
					lastTrafBox1 = moofBox->trafBox1;
					lastTrafBox2 = moofBox->trafBox2;
					delete moofBox;
					moofBox = nullptr;
				}
				moofBox = new MoofBox(box);
				box = nullptr;
				if (!moofBox->parse()) {
					delete moofBox;
					break;
				}
				if (lastTrafBox1.trunBox.sampleCount > 0) {
					moofBox->trafBox1.trunBox.eachSampleDuration =
						(moofBox->trafBox1.tfdtBox.baseMediaDecodeTime - lastTrafBox1.tfdtBox.baseMediaDecodeTime)
						/ lastTrafBox1.trunBox.sampleCount;
				}
				if (lastTrafBox2.trunBox.sampleCount > 0) {
					moofBox->trafBox2.trunBox.eachSampleDuration =
						(moofBox->trafBox2.tfdtBox.baseMediaDecodeTime - lastTrafBox2.tfdtBox.baseMediaDecodeTime)
						/ lastTrafBox2.trunBox.sampleCount;
				}
			} else if (strcmp(box->type, "mfra") == 0) {
				break;
			}
			if (nullptr != moofBox) {
				moofBox->moofOffset = offset;
			}
		} while (true);
	} while (false);
	if (nullptr != fp) {
		fp->close();
		delete fp;
		fp = nullptr;
	}
	return moofBox;
}

MfraBox* CFMP4Appender::getMfraBox(const char* filePath)
{
	unsigned long long offset = 0;
	std::fstream* fp = nullptr;
	MfraBox* mfraBox = nullptr;
	do
	{
		fp = new std::fstream(filePath, std::ios::binary | std::ios::in);
		if (!fp->is_open()) {
			break;
		}

		Box* box = nullptr;
		do
		{
			if (nullptr != box) {
				delete box;
				box = nullptr;
			}
			box = new Box;
			if (!parseBox(fp, box, true)) {
				break;
			}
			if (strcmp(box->type, "mfra") == 0) {
				mfraBox = new MfraBox(box);
				if (!mfraBox->parse()) {
					delete mfraBox;
					mfraBox = nullptr;
				}
				box = nullptr;
				break;
			}
		} while (true);
		if (nullptr != box) {
			delete box;
			box = nullptr;
		}
	} while (false);
	if (nullptr != fp) {
		fp->close();
		delete fp;
		fp = nullptr;
	}
	return mfraBox;
}

int CFMP4Appender::ftruncateFile(const char* filePath, int length)
{
	int ret = -1;
	do 
	{
		std::fstream file(filePath, std::ios::binary | std::ios::in | std::ios::out);
		if (!file.is_open()) {
			printf("Failed to open file\n");
			break;
		}
		file.seekp(length);
		file.write("", 1);
		file.flush();
		file.seekp(0, std::ios::end);
		ret = file.tellp();
		file.close();
	} while (false);
	return ret;
}

bool TfdtBox::parse()
{
	bool ret = false;
	do 
	{
		if (size < 12) {
			printf("ERROR: Invalid tfdt box size.\n");
			break;
		}

		version = buffer[8];

		if (version == 0) {
			if (size < 16) {
				printf("ERROR: Invalid tfdt box size for version 0.\n");
				break;
			}
			unsigned int time = 0;
	/*		memcpy(&time, buffer + 12, sizeof(unsigned int));
			std::reverse((char*)&time, (char*)&time + sizeof(unsigned int));*/
			time = read32(buffer + 12);
			baseMediaDecodeTime = time;
			//baseMediaDecodeTime = (buffer[12] << 24) | (buffer[13] << 16) | (buffer[14] << 8) | buffer[15];
		}
		else if (version == 1) {
			if (size < 20) {
				printf("ERROR: Invalid tfdt box size for version 1.\n");
				break;
			}

			//memcpy(&baseMediaDecodeTime, buffer + 12, sizeof(unsigned long long));
			//std::reverse((char*)&baseMediaDecodeTime, (char*)&baseMediaDecodeTime + sizeof(unsigned long long));
			baseMediaDecodeTime = read64(buffer + 12);
		}
		else {
			printf("ERROR: Unsupported tfdt box version.\n");
			break;
		}
		ret = true;
	} while (false);
	return ret;
}

bool TfdtBox::modifyBaseMediaDecodeTime(unsigned long long timeBase)
{
	bool ret = false;
	do
	{
		if (nullptr == buffer) {
			break;
		}
		if (size < 12) {
			printf("ERROR: Invalid tfdt box size.\n");
			break;
		}
		if (version == 0) {
			if (size < 16) {
				printf("ERROR: Invalid tfdt box size for version 0.\n");
				break;
			}
			unsigned int time = timeBase + baseMediaDecodeTime;
			std::reverse((char*)&time, (char*)&time + sizeof(unsigned int));
			memcpy(buffer + 12, &time, sizeof(unsigned int));
		}
		else if (version == 1) {
			if (size < 20) {
				printf("ERROR: Invalid tfdt box size for version 1.\n");
				break;
			}
			unsigned long long time = timeBase + baseMediaDecodeTime;
			std::reverse((char*)&time, (char*)&time + sizeof(unsigned long long));
			memcpy(buffer + 12, &time, sizeof(unsigned long long));
		}
		else {
			printf("ERROR: Unsupported tfdt box version.\n");
			break;
		}
		ret = true;
	} while (false);
	return ret;
}

bool MfhdBox::parse()
{
	bool ret = false;
	do 
	{
		if (size < 12) {
			printf("ERROR: Invalid mfhd box size.\n");
			break;
		}

		version = buffer[8];

		if (version == 0) {
			if (size < 16) {
				printf("ERROR: Invalid mfhd box size for version 0.\n");
				break;
			}
			sequenceNumber = read32(buffer + 12);
		}
		else {
			printf("ERROR: Unsupported mfhd box version.\n");
			break;
		}
		ret = true;
	} while (false);
	return ret;
}

bool MfhdBox::modifySequenceNumber(unsigned int baseNumber)
{
	bool ret = false;
	do
	{
		if (nullptr == buffer) {
			break;
		}
		if (size < 12) {
			printf("ERROR: Invalid mfhd box size.\n");
			break;
		}

		if (version == 0) {
			if (size < 16) {
				printf("ERROR: Invalid mfhd box size for version 0.\n");
				break;
			}
			unsigned int number = baseNumber + sequenceNumber;
			std::reverse((char*)&number, (char*)&number + sizeof(unsigned int));
			memcpy(buffer + 12, &number, sizeof(unsigned int));
		}
		else {
			printf("ERROR: Unsupported mfhd box version.\n");
			break;
		}
		ret = true;
	} while (false);
	return ret;
}

bool TfhdBox::parse()
{
	bool ret = false;
	do 
	{
		if (size < 16) {
			printf("ERROR: Invalid tfhd box size.\n");
			break;
		}

		unsigned long long offset = 8;
		version = buffer[8];
		flags = (buffer[9] << 16) | (buffer[10] << 8) | buffer[11];
		offset += 4;

		if (flags & 0b01) {
			//trackId = (buffer[12] << 24) | (buffer[13] << 16) | (buffer[14] << 8) | buffer[15];
			trackId = read32(buffer + offset);
			offset += sizeof(unsigned int);

			if (flags & 0b1) {
				baseDataOffset = read64(buffer + offset);
				offset += sizeof(unsigned long long);
			}

			if (flags & 0b10) {
				sampleDescriptionIndex = read32(buffer + offset);
				offset += sizeof(unsigned int);
			}

			if (flags & 0b1000) {
				sampleDuration = read32(buffer + offset);
				offset += sizeof(unsigned int);
			}

			if (flags & 0b10000) {
				sampleSize = read32(buffer + offset);
				offset += sizeof(unsigned int);
			}
		}
		else {
			printf("ERROR: tfhd box missing track ID flag.\n");
			break;
		}
		ret = true;
	} while (false);
	return ret;
}

bool TfhdBox::modifyBaseDataOffset(unsigned long long baseOffset)
{
	bool ret = false;
	do
	{
		if (size < 16) {
			printf("ERROR: Invalid tfhd box size.\n");
			break;
		}

		unsigned long long offset = 8;
		offset += 4;

		if (flags & 0b01) {
			offset += sizeof(unsigned int);

			if (flags & 0b1) {
				std::reverse((char*)&baseOffset, (char*)&baseOffset + sizeof(unsigned long long));
				memcpy(buffer + offset, &baseOffset, sizeof(unsigned long long));
				offset += sizeof(unsigned long long);
			}
		}
		else {
			printf("ERROR: tfhd box missing track ID flag.\n");
			break;
		}
		ret = true;
	} while (false);
	return ret;
}

bool TrunBox::parse()
{
	bool ret = false;
	do
	{
		if (size < 16) {
			printf("ERROR: Invalid trun box size.\n");
			break;
		}

		unsigned long long offset = 8;
		version = buffer[8];
		flags = (buffer[9] << 16) | (buffer[10] << 8) | buffer[11];
		offset += 4;
		sampleCount = read32(buffer + offset);
		ret = true;
	} while (false);
	return ret;
}

bool TrafBox::parse()
{
	bool ret = false;
	do 
	{
		parseBox(buffer + 8, &tfhdBox);
		if (strcmp(tfhdBox.type, "tfhd") != 0) {
			break;
		}
		if (!tfhdBox.parse()) {
			break;
		}
		parseBox(buffer + 8 + tfhdBox.size, &tfdtBox);
		if (strcmp(tfdtBox.type, "tfdt") != 0) {
			break;
		}
		if (!tfdtBox.parse()) {
			break;
		}
		parseBox(buffer + 8 + tfhdBox.size + tfdtBox.size, &trunBox);
		if (strcmp(trunBox.type, "trun") != 0) {
			break;
		}
		if (!trunBox.parse()) {
			break;
		}
		ret = true;
	} while (false);
	return ret;
}

MoofBox::MoofBox(Box* box)
{
	moofBox = box;
}

MoofBox::~MoofBox()
{
	if (nullptr != moofBox) {
		if (moofBox->buffer) {
			delete[] moofBox->buffer;
			moofBox->buffer = nullptr;
		}
		delete moofBox;
	}
}

bool MoofBox::parse()
{
	bool ret = false;
	do 
	{
		parseBox(moofBox->buffer + 8, &mfhdBox);
		if (strcmp(mfhdBox.type, "mfhd") != 0) {
			break;
		}
		if (!mfhdBox.parse()) {
			break;
		}
		parseBox(moofBox->buffer + 8 + mfhdBox.size, &trafBox1);
		if (strcmp(trafBox1.type, "traf") != 0) {
			break;
		}
		if (!trafBox1.parse()) {
			break;
		}
		if (mfhdBox.size + trafBox1.size + 8 < moofBox->size) {
			parseBox(moofBox->buffer + (8 + mfhdBox.size + trafBox1.size), &trafBox2);
			if (strcmp(trafBox2.type, "traf") != 0) {
				break;
			}
			if (!trafBox2.parse()) {
				break;
			}
		}
		ret = true;
	} while (false);
	return ret;
}

bool TfraBox::parse()
{
	bool ret = false;
	do
	{
		unsigned long long offset = 8;
		version = buffer[offset];
		offset += 4;
		trackId = read32(buffer + offset);
		offset += sizeof(unsigned int);

		offset = size - 3 - 16;
		time = read64(buffer + offset);
		offset += 8;
		moofOffset = read64(buffer + offset);
		ret = true;
	} while (false);
	return ret;
}

bool TfraBox::modifyTime(unsigned long long baseTime)
{
	bool ret = false;
	do
	{
		unsigned long long offset = size - 3 - 16;
		std::reverse((char*)&baseTime, (char*)&baseTime + sizeof(unsigned long long));
		memcpy(buffer + offset, &baseTime, sizeof(unsigned long long));
		ret = true;
	} while (false);
	return ret;
}

bool TfraBox::modifyMoofOffset(unsigned long long baseMoofOffset)
{
	bool ret = false;
	do
	{
		unsigned long long offset = size - 3 - 8;
		std::reverse((char*)&baseMoofOffset, (char*)&baseMoofOffset + sizeof(unsigned long long));
		memcpy(buffer + offset, &baseMoofOffset, sizeof(unsigned long long));
		ret = true;
	} while (false);
	return ret;
}

MfraBox::MfraBox(Box* box)
{
	mfraBox = box;
}

MfraBox::~MfraBox()
{
	if (nullptr != mfraBox) {
		if (mfraBox->buffer) {
			delete[] mfraBox->buffer;
			mfraBox->buffer = nullptr;
		}
		delete mfraBox;
	}
}

bool MfraBox::parse()
{
	bool ret = false;
	do 
	{
		if (!parseBox(mfraBox->buffer + 8, &tfraBox1)) {
			break;
		}
		if (!tfraBox1.parse()) {
			break;
		}
		if (tfraBox1.size * 2 + 8 > mfraBox->size) {
			break;
		}
		if (!parseBox(mfraBox->buffer + (8 + tfraBox1.size), &tfraBox2)) {
			break;
		}
		if (!tfraBox2.parse()) {
			break;
		}
		ret = true;
	} while (false);
	return ret;
}
